import os
import matplotlib.pyplot as plt
from mpl_toolkits import mplot3d
import numpy as np
import time

def benchmark(command):
    start_time = time.time()
    os.system(command)
    end_time = time.time()
    performance_metric = end_time - start_time
    return performance_metric

sizes = [256, 512, 1024, 2048, 4096]
blocks = [1, 2, 4, 8]
cores = [1, 2, 4, 8]
results = []

for size in sizes:
    performance_metrics = []
    for block in blocks:
        for core in cores:
            command = f'mpiexec -n {core} python3 main.py {size} {block}'
            performance_metric = benchmark(command)
            performance_metrics.append(performance_metric)
    results.append(performance_metrics)

# 生成3D图表
fig = plt.figure(figsize=(10, 6))
ax = plt.axes(projection='3d')

x = []
y = []
z = []
for i, size in enumerate(sizes):
    for j, block in enumerate(blocks):
        for k, core in enumerate(cores):
            x.append(i)
            y.append(j)
            z.append(k)

dx = dy = 0.8
dz = np.ravel(results)  # 将结果列表展平为一维数组

# 绘制3D图形
ax.bar3d(x, y, z, dx, dy, dz)

ax.set_xlabel('Matrix Size')
ax.set_ylabel('Blocks')
ax.set_zlabel('Number of Processes')
ax.set_title('Performance Comparison for Different Matrix Sizes, Blocks, and Number of Processes')

plt.xticks(np.arange(len(sizes)), sizes)
plt.yticks(np.arange(len(blocks)), blocks)
ax.set_zticks(np.arange(len(cores)))
ax.set_zticklabels(cores)

plt.show()